package com.google.code.sfa.util;

import org.biojava.bio.SimpleAnnotation;
import org.biojava.bio.seq.impl.SimpleSequence;
import org.biojava.bio.seq.impl.SimpleGappedSequence;
import org.biojava.bio.seq.Sequence;
import org.biojava.bio.symbol.SymbolList;
import org.biojava.bio.symbol.Symbol;
import org.biojava.bio.symbol.IllegalSymbolException;
import org.biojava.bio.seq.DNATools;
import org.biojava.bio.seq.io.SymbolTokenization;
import org.biojava.bio.BioException;
import org.biojava.bio.BioRuntimeException;
import org.biojava.bio.symbol.Alignment;
import org.biojava.bio.symbol.SimpleAlignment;
import org.biojava.bio.symbol.SimpleSymbolList;

import java.util.Map;
import java.util.HashMap;

public class OverlapAlignment {

    private int[][] matrix;
    private Sequence query, subject;
    private int delete, insert, match, replace;
    public int score;
    public int offset;
    public int reverseOffset;
    public Alignment alignment;

    public OverlapAlignment() {
	insert = 4;
	delete = 4;
	match = 5;
	replace = -4;
    }

    public void align( SymbolList query_list, SymbolList subject_list) {
	SymbolTokenization st;
	try {
	    st = DNATools.getDNA().getTokenization("default");
	} catch( BioException e) {
	    throw new BioRuntimeException(e);
	}

	query = makeSequence( query_list, "query");
	subject = makeSequence( subject_list, "subject");
	
	matrix = new int[query.length() + 1][subject.length() + 1];

	Position pos = null;
	while( (pos = next(pos)) != null) {
	    matrix[pos.row][pos.col] = score( pos);
	}
	//printMatrix();
	String queryAlign = "";
	String subjectAlign = "";
	pos = null;
	Position last = null;
	while( (pos = trace(pos)) != null) {
	    //System.out.println( pos.toString());
	    if( last == null ) {
		score = matrix[ pos.row][ pos.col];
		if( pos.row == subject.length() ) {
		    reverseOffset = subject.length() - pos.col + 1;
		} else {
		    reverseOffset = -(query.length() - pos.row) + 1;
		}
	    } else {
	    
	    try {
		if( pos.row < last.row ) {
		    queryAlign = st.tokenizeSymbol(query.symbolAt(last.row)) +
			queryAlign;
		} else {
		    queryAlign = '-' + queryAlign;
		}
	    } catch( IllegalSymbolException e) {
		queryAlign = '?' + queryAlign;
	    }
	    try {
		if( pos.col < last.col) {
		    subjectAlign = st.tokenizeSymbol(subject.symbolAt(last.col)) + 
			subjectAlign;
		} else {
		    subjectAlign = '-' + subjectAlign;
		}
	    } catch( IllegalSymbolException e) {
		subjectAlign = '?' + subjectAlign;
	    }
	    }
	    last = pos;
	}
	   
	//System.out.println( subjectAlign);
	//System.out.println( queryAlign);

	if( last != null) {
	    if( last.row == 0 ) {
		offset = last.col + 1;
	    } else {
		offset = -last.row + 1;
	    }
	} else {
	    offset = 0;
	}

	

	try {
	query = new SimpleGappedSequence(
	    new SimpleSequence(
	        new SimpleSymbolList( query.getAlphabet().getTokenization("token"),
				      queryAlign),
		query.getURN(), query.getName(), query.getAnnotation()));
	subject = new SimpleGappedSequence(
	    new SimpleSequence(
	        new SimpleSymbolList( subject.getAlphabet().getTokenization("token"),
				      subjectAlign),
		subject.getURN(), subject.getName(), subject.getAnnotation()));
	Map<String, Sequence> m = new HashMap<String, Sequence>();
	m.put( query.getName(), query);
	m.put( subject.getName(), subject);
	alignment = new SimpleAlignment( m);

	} catch( Exception e) {
	    throw new BioRuntimeException( e);
	}
    }

    private int score( Position pos) {
	if( pos.row == 0 || pos.col == 0) {
	    return 0;
	}
	/*if( pos.row == 0) {
	    if( pos.col == 0) {
		return 0;
	    } else {
		return matrix[0][ pos.col - 1] + insert;
	    }
	} else if( pos.col == 0) {
	    return matrix[ pos.row - 1][0] + delete;
	}
	*/

	return max(
		   matrix[pos.row - 1][pos.col - 1] + score( query.symbolAt(pos.row),
							     subject.symbolAt(pos.col)),
		   matrix[pos.row - 1][pos.col] - delete,
		   matrix[pos.row][pos.col - 1] - insert);
    }

    private Position trace( Position pos) {
	if( pos == null) {
	    int row, col;
	    col = subject.length();
	    int maxValue = matrix[0][col];
	    int maxRow = 0;
	    int maxCol = col;
	    for( row = 0; row <= query.length(); row++) {
		if( matrix[row][col] > maxValue) {
		    maxValue = matrix[row][col];
		    maxRow = row;
		    maxCol = col;
		}
	    }
	    row = query.length();
	    for( col = 0; col <= subject.length(); col++) {
		if( matrix[row][col] > maxValue) {
		    maxValue = matrix[row][col];
		    maxRow = row;
		    maxCol = col;
		}
	    }
	    if( maxRow == 0 || maxCol == 0) {
		return null;
	    }
	    return new Position(maxRow, maxCol);
	} else if( pos.col == 0 || pos.row == 0) {
	    return null;
	} else if( matrix[pos.row][pos.col] == 
	    matrix[pos.row-1][pos.col-1] + score( query.symbolAt(pos.row),
						  subject.symbolAt(pos.col)) ) {
	    return new Position( pos.row - 1, pos.col - 1);
	} else if( matrix[pos.row][pos.col] == matrix[pos.row][pos.col - 1] - insert) {
	    return new Position( pos.row, pos.col - 1);
	} else if( matrix[pos.row][pos.col] == matrix[pos.row-1][pos.col] - delete) {
	    return new Position( pos.row - 1, pos.col);
	} else {
	    System.out.println("Error on traceback");
	    return null;
	}
    }

    private int score( Symbol q, Symbol s) {
	//if( q.getMatches().contains( s) || s.getMatches().contains( q)) {
	if( q == s) {
	    return match;
	} else {
	    return replace;
	}
    }

    private Position next( Position last) {
	if( last == null) {
	    return new Position( 0, 0);
	}
	if( last.row == query.length()) {
	    if( last.col == subject.length() ) {
		return null;
	    } else {
		return new Position( 0, last.col + 1);
	    }
	} else {
	    return new Position( last.row + 1, last.col);
	}
    }

    private Sequence makeSequence( SymbolList list, String name) {
	if( list instanceof Sequence) {
	    return (Sequence) list;
	} else {
	    return new SimpleSequence( list, "", name, new SimpleAnnotation() );
	}
    }

  protected int max(int x, int y, int z) {
    if ((x > y) && (x > z))
      return x;
    if (y > z)
      return y;
    return z;
  }

    private void printMatrix() {
	try {
	    System.out.print("     ");
	    for( int col = 1; col <= subject.length(); col++) {
		System.out.print( "   " + DNATools.dnaToken( subject.symbolAt( col)));
	    }
	    System.out.println("");
	    for( int row = 0; row <= query.length(); row++) {
		if( row == 0) {
		    System.out.print(" ");
		} else {
		    System.out.print( DNATools.dnaToken( query.symbolAt( row)));
		}
		for( int col = 0; col <= subject.length(); col++) {
		    System.out.printf( " %3d", matrix[row][col]);
		}
		System.out.println("");
	    }
	} catch( Exception e) {
	    e.printStackTrace();
	}
    }

}

class Position {
    public int row;
    public int col;

    public Position( int row, int col) {
	this.row = row;
	this.col = col;
    }

    public String toString() {
	return "(" + row + "," + col + ")";
    }
}